"""
基于 CrewAI 的数据收集智能体
"""
from crewai import Task
from agents.base_agent import CrewAIStockAgent
from tools.stock_tools import StockDataTool, NewsSearchTool
from typing import List, Dict, Any
from colorama import Fore, Style


class DataCollectorAgent(CrewAIStockAgent):
    """基于 CrewAI 的数据收集智能体"""
    
    def __init__(self):
        # 创建工具实例
        tools = [
            StockDataTool(),
            NewsSearchTool()
        ]
        super().__init__("data_collector", tools)
        print(Fore.CYAN + f"    - [CrewAI] {self.config.name} 已配置 {len(tools)} 个工具")
    
    def create_tasks(self, stock_code: str) -> List[Task]:
        """创建数据收集任务"""
        tasks = []
        
        # 任务1：收集股票历史数据
        stock_data_task = Task(
            description=f"""
            收集股票 {stock_code} 的历史数据：
            1. 获取最近一年的股票价格数据
            2. 分析价格走势和成交量变化
            3. 计算关键统计指标
            4. 整理数据为结构化报告
            
            请使用股票数据工具获取准确的数据。
            """,
            expected_output="""
            输出应包含：
            - 股票基本信息
            - 最新价格数据
            - 历史价格统计
            - 成交量分析
            - 数据获取时间
            """,
            agent=self.agent,
            tools=self.tools
        )
        tasks.append(stock_data_task)
        
        # 任务2：收集相关新闻
        news_task = Task(
            description=f"""
            搜索股票 {stock_code} 的相关新闻和市场资讯：
            1. 搜索最新的财经新闻
            2. 获取公司公告和重要事件
            3. 分析市场情绪和舆论导向
            4. 整理新闻摘要
            
            请使用新闻搜索工具获取最新资讯。
            """,
            expected_output="""
            输出应包含：
            - 最新新闻列表（至少5条）
            - 新闻时间和来源
            - 内容摘要
            - 市场情绪分析
            """,
            agent=self.agent,
            tools=self.tools
        )
        tasks.append(news_task)
        
        # 任务3：整合数据报告
        integration_task = Task(
            description=f"""
            整合股票 {stock_code} 的所有收集到的数据：
            1. 将股票数据和新闻信息整合
            2. 创建结构化的数据收集报告
            3. 标注数据来源和时效性
            4. 提供数据质量评估
            
            基于前面收集的数据，生成完整的数据收集报告。
            """,
            expected_output="""
            完整的数据收集报告，包括：
            - 执行摘要
            - 股票数据部分
            - 新闻资讯部分
            - 数据质量说明
            - 收集时间戳
            """,
            agent=self.agent,
            context=[stock_data_task, news_task]  # 依赖前两个任务
        )
        tasks.append(integration_task)
        
        print(Fore.GREEN + f"    - [CrewAI] {self.config.name} 创建了 {len(tasks)} 个任务")
        return tasks
    
    def get_summary(self) -> str:
        """获取智能体摘要"""
        return f"""
        📊 数据收集智能体
        角色: {self.config.role}
        目标: {self.config.goal}
        工具: {', '.join([tool.name for tool in self.tools])}
        """
